#!/usr/bin/env python3
"""
Vision Service Router

Implements intelligent routing and fallback mechanisms for vision services.
Handles service availability detection, load balancing, and automatic failover
between Coze API and local vision models.

Features:
- Service availability monitoring
- Intelligent service selection based on performance and availability
- Automatic failover and recovery
- Load balancing and rate limiting
- Service health scoring
- Circuit breaker pattern implementation
"""

import asyncio
import time
import logging
from typing import Optional, Dict, Any, List, Callable, Tuple
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum
import statistics
import json

from .coze_integration import CozeIntegration, VisionQueryRequest, VisionQueryResponse, CozeIntegrationError
from .local_vision_model import LocalVisionModel


class ServiceStatus(Enum):
    """Service status states"""
    HEALTHY = "healthy"
    DEGRADED = "degraded"
    UNHEALTHY = "unhealthy"
    CIRCUIT_OPEN = "circuit_open"
    UNKNOWN = "unknown"


class ServiceType(Enum):
    """Available service types"""
    COZE_API = "coze_api"
    LOCAL_MODEL = "local_model"


@dataclass
class ServiceMetrics:
    """Service performance metrics"""
    total_requests: int = 0
    successful_requests: int = 0
    failed_requests: int = 0
    average_response_time: float = 0.0
    last_success_time: Optional[datetime] = None
    last_failure_time: Optional[datetime] = None
    consecutive_failures: int = 0
    consecutive_successes: int = 0
    response_times: List[float] = field(default_factory=list)
    error_rate: float = 0.0
    availability: float = 1.0
    
    def update_success(self, response_time: float):
        """Update metrics for successful request"""
        self.total_requests += 1
        self.successful_requests += 1
        self.consecutive_successes += 1
        self.consecutive_failures = 0
        self.last_success_time = datetime.now()
        
        # Update response times (keep last 100)
        self.response_times.append(response_time)
        if len(self.response_times) > 100:
            self.response_times.pop(0)
        
        # Update average response time
        self.average_response_time = statistics.mean(self.response_times)
        
        # Update error rate and availability
        self._update_rates()
    
    def update_failure(self):
        """Update metrics for failed request"""
        self.total_requests += 1
        self.failed_requests += 1
        self.consecutive_failures += 1
        self.consecutive_successes = 0
        self.last_failure_time = datetime.now()
        
        # Update error rate and availability
        self._update_rates()
    
    def _update_rates(self):
        """Update error rate and availability"""
        if self.total_requests > 0:
            self.error_rate = self.failed_requests / self.total_requests
            self.availability = self.successful_requests / self.total_requests
        else:
            self.error_rate = 0.0
            self.availability = 1.0
    
    def get_health_score(self) -> float:
        """Calculate health score (0.0 to 1.0)"""
        if self.total_requests == 0:
            return 1.0
        
        # Base score from availability
        score = self.availability
        
        # Penalty for consecutive failures
        if self.consecutive_failures > 0:
            failure_penalty = min(0.5, self.consecutive_failures * 0.1)
            score -= failure_penalty
        
        # Bonus for consecutive successes
        if self.consecutive_successes > 5:
            success_bonus = min(0.2, (self.consecutive_successes - 5) * 0.02)
            score += success_bonus
        
        # Response time factor
        if self.average_response_time > 0:
            # Penalty for slow responses (>10s gets penalty)
            if self.average_response_time > 10.0:
                time_penalty = min(0.3, (self.average_response_time - 10.0) * 0.01)
                score -= time_penalty
        
        return max(0.0, min(1.0, score))


@dataclass
class CircuitBreakerState:
    """Circuit breaker state for a service"""
    is_open: bool = False
    failure_count: int = 0
    last_failure_time: Optional[datetime] = None
    next_attempt_time: Optional[datetime] = None
    failure_threshold: int = 5
    timeout_duration: int = 60  # seconds
    
    def should_attempt_request(self) -> bool:
        """Check if request should be attempted"""
        if not self.is_open:
            return True
        
        if self.next_attempt_time and datetime.now() >= self.next_attempt_time:
            return True
        
        return False
    
    def record_success(self):
        """Record successful request"""
        self.is_open = False
        self.failure_count = 0
        self.last_failure_time = None
        self.next_attempt_time = None
    
    def record_failure(self):
        """Record failed request"""
        self.failure_count += 1
        self.last_failure_time = datetime.now()
        
        if self.failure_count >= self.failure_threshold:
            self.is_open = True
            self.next_attempt_time = datetime.now() + timedelta(seconds=self.timeout_duration)


class VisionServiceRouter:
    """
    Intelligent router for vision services with fallback mechanisms.
    
    Manages service selection, health monitoring, and automatic failover
    between different vision service providers.
    """
    
    def __init__(self, logger=None):
        """
        Initialize the vision service router.
        
        Args:
            logger: Logger instance
        """
        self.logger = logger or logging.getLogger(__name__)
        
        # Service instances
        self.services: Dict[ServiceType, Any] = {}
        self.service_metrics: Dict[ServiceType, ServiceMetrics] = {}
        self.circuit_breakers: Dict[ServiceType, CircuitBreakerState] = {}
        
        # Router configuration
        self.preferred_service = ServiceType.COZE_API
        self.enable_fallback = True
        self.enable_load_balancing = False
        self.health_check_interval = 30.0  # seconds
        self.metrics_window_size = 100
        
        # Service weights for load balancing
        self.service_weights: Dict[ServiceType, float] = {
            ServiceType.COZE_API: 1.0,
            ServiceType.LOCAL_MODEL: 0.8
        }
        
        # Rate limiting
        self.rate_limits: Dict[ServiceType, Dict[str, Any]] = {
            ServiceType.COZE_API: {
                'max_requests_per_minute': 60,
                'current_requests': 0,
                'window_start': datetime.now()
            },
            ServiceType.LOCAL_MODEL: {
                'max_requests_per_minute': 120,
                'current_requests': 0,
                'window_start': datetime.now()
            }
        }
        
        # Initialize metrics and circuit breakers
        for service_type in ServiceType:
            self.service_metrics[service_type] = ServiceMetrics()
            self.circuit_breakers[service_type] = CircuitBreakerState()
        
        # Background tasks
        self._health_check_task = None
        self._start_background_tasks()
    
    def register_service(self, service_type: ServiceType, service_instance: Any):
        """
        Register a vision service instance.
        
        Args:
            service_type: Type of service
            service_instance: Service instance (CozeIntegration or LocalVisionModel)
        """
        self.services[service_type] = service_instance
        self.logger.info(f"Registered service: {service_type.value}")
    
    def set_preferred_service(self, service_type: ServiceType):
        """Set preferred service type."""
        self.preferred_service = service_type
        self.logger.info(f"Preferred service set to: {service_type.value}")
    
    def enable_service_fallback(self, enable: bool):
        """Enable or disable service fallback."""
        self.enable_fallback = enable
        self.logger.info(f"Service fallback {'enabled' if enable else 'disabled'}")
    
    async def route_request(self, request: VisionQueryRequest) -> VisionQueryResponse:
        """
        Route vision request to the best available service.
        
        Args:
            request: Vision query request
            
        Returns:
            Vision query response
        """
        # Get ordered list of services to try
        service_order = self._get_service_order()
        
        last_error = None
        
        for service_type in service_order:
            # Check if service is available
            if not self._is_service_available(service_type):
                continue
            
            # Check circuit breaker
            if not self.circuit_breakers[service_type].should_attempt_request():
                self.logger.warning(f"Circuit breaker open for {service_type.value}")
                continue
            
            # Check rate limiting
            if not self._check_rate_limit(service_type):
                self.logger.warning(f"Rate limit exceeded for {service_type.value}")
                continue
            
            # Attempt request
            try:
                start_time = time.time()
                response = await self._execute_request(service_type, request)
                response_time = time.time() - start_time
                
                # Record success
                self._record_success(service_type, response_time)
                
                # Update response metadata
                response.method_used = f"{service_type.value}_routed"
                
                return response
                
            except Exception as e:
                last_error = e
                self.logger.error(f"Service {service_type.value} failed: {e}")
                
                # Record failure
                self._record_failure(service_type)
                
                # Continue to next service if fallback is enabled
                if not self.enable_fallback:
                    break
        
        # All services failed
        return self._create_error_response(request, last_error)
    
    def _get_service_order(self) -> List[ServiceType]:
        """Get ordered list of services to try."""
        if not self.enable_load_balancing:
            # Simple preferred service first, then others
            services = [self.preferred_service]
            for service_type in ServiceType:
                if service_type != self.preferred_service:
                    services.append(service_type)
            return services
        else:
            # Load balancing based on health scores and weights
            return self._get_load_balanced_order()
    
    def _get_load_balanced_order(self) -> List[ServiceType]:
        """Get service order based on load balancing algorithm."""
        service_scores = []
        
        for service_type in ServiceType:
            if not self._is_service_available(service_type):
                continue
            
            # Calculate composite score
            health_score = self.service_metrics[service_type].get_health_score()
            weight = self.service_weights.get(service_type, 1.0)
            circuit_penalty = 0.5 if self.circuit_breakers[service_type].is_open else 0.0
            
            composite_score = (health_score * weight) - circuit_penalty
            service_scores.append((service_type, composite_score))
        
        # Sort by score (highest first)
        service_scores.sort(key=lambda x: x[1], reverse=True)
        
        return [service_type for service_type, _ in service_scores]
    
    def _is_service_available(self, service_type: ServiceType) -> bool:
        """Check if service is available."""
        if service_type not in self.services:
            return False
        
        service = self.services[service_type]
        
        if service_type == ServiceType.COZE_API:
            return hasattr(service, 'is_initialized') and service.is_initialized
        elif service_type == ServiceType.LOCAL_MODEL:
            return hasattr(service, 'is_model_loaded') and service.is_model_loaded()
        
        return False
    
    def _check_rate_limit(self, service_type: ServiceType) -> bool:
        """Check if service is within rate limits."""
        if service_type not in self.rate_limits:
            return True
        
        rate_limit = self.rate_limits[service_type]
        current_time = datetime.now()
        
        # Reset window if needed
        if (current_time - rate_limit['window_start']).total_seconds() >= 60:
            rate_limit['current_requests'] = 0
            rate_limit['window_start'] = current_time
        
        # Check limit
        if rate_limit['current_requests'] >= rate_limit['max_requests_per_minute']:
            return False
        
        # Increment counter
        rate_limit['current_requests'] += 1
        return True
    
    async def _execute_request(self, service_type: ServiceType, request: VisionQueryRequest) -> VisionQueryResponse:
        """Execute request on specific service."""
        service = self.services[service_type]
        
        if service_type == ServiceType.COZE_API:
            return await service.process_vision_query(request)
        elif service_type == ServiceType.LOCAL_MODEL:
            # Convert request for local model
            from .coze_integration import convert_image_array_to_bytes
            import cv2
            import numpy as np
            
            # Convert bytes to image
            nparr = np.frombuffer(request.image_data, np.uint8)
            image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            
            # Process with local model
            result_text = service.analyze_image(image, request.user_query)
            
            # Create response
            return VisionQueryResponse(
                query_id=request.query_id,
                response_text=result_text,
                confidence_score=0.8,
                processing_time_ms=0,  # Will be calculated by caller
                model_used="local_model",
                method_used="local",
                timestamp=datetime.now()
            )
        
        raise ValueError(f"Unknown service type: {service_type}")
    
    def _record_success(self, service_type: ServiceType, response_time: float):
        """Record successful request."""
        self.service_metrics[service_type].update_success(response_time)
        self.circuit_breakers[service_type].record_success()
        
        self.logger.debug(f"Success recorded for {service_type.value}: {response_time:.2f}s")
    
    def _record_failure(self, service_type: ServiceType):
        """Record failed request."""
        self.service_metrics[service_type].update_failure()
        self.circuit_breakers[service_type].record_failure()
        
        self.logger.debug(f"Failure recorded for {service_type.value}")
    
    def _create_error_response(self, request: VisionQueryRequest, error: Optional[Exception]) -> VisionQueryResponse:
        """Create error response when all services fail."""
        error_message = "所有视觉服务都不可用"
        if error:
            error_message += f": {str(error)}"
        
        return VisionQueryResponse(
            query_id=request.query_id,
            response_text=error_message,
            confidence_score=0.0,
            processing_time_ms=0,
            model_used="none",
            method_used="error",
            timestamp=datetime.now(),
            error=error_message
        )
    
    def _start_background_tasks(self):
        """Start background maintenance tasks."""
        self._health_check_task = asyncio.create_task(self._health_check_loop())
    
    async def _health_check_loop(self):
        """Background health check loop."""
        while True:
            try:
                await asyncio.sleep(self.health_check_interval)
                await self._perform_health_checks()
            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error(f"Health check failed: {e}")
    
    async def _perform_health_checks(self):
        """Perform health checks on all services."""
        for service_type, service in self.services.items():
            try:
                is_healthy = await self._check_service_health(service_type, service)
                
                if is_healthy:
                    # Reset circuit breaker if service is healthy
                    if self.circuit_breakers[service_type].is_open:
                        self.circuit_breakers[service_type].record_success()
                        self.logger.info(f"Service {service_type.value} recovered")
                
            except Exception as e:
                self.logger.error(f"Health check failed for {service_type.value}: {e}")
    
    async def _check_service_health(self, service_type: ServiceType, service: Any) -> bool:
        """Check health of specific service."""
        try:
            if service_type == ServiceType.COZE_API:
                if hasattr(service, 'test_connection'):
                    result = await service.test_connection()
                    return result.get('overall', False)
                return hasattr(service, 'is_initialized') and service.is_initialized
            
            elif service_type == ServiceType.LOCAL_MODEL:
                return hasattr(service, 'is_model_loaded') and service.is_model_loaded()
            
            return False
            
        except Exception as e:
            self.logger.error(f"Health check error for {service_type.value}: {e}")
            return False
    
    def get_service_status(self) -> Dict[str, Any]:
        """Get current status of all services."""
        status = {}
        
        for service_type in ServiceType:
            metrics = self.service_metrics[service_type]
            circuit_breaker = self.circuit_breakers[service_type]
            
            # Determine service status
            if not self._is_service_available(service_type):
                service_status = ServiceStatus.UNKNOWN
            elif circuit_breaker.is_open:
                service_status = ServiceStatus.CIRCUIT_OPEN
            elif metrics.get_health_score() > 0.8:
                service_status = ServiceStatus.HEALTHY
            elif metrics.get_health_score() > 0.5:
                service_status = ServiceStatus.DEGRADED
            else:
                service_status = ServiceStatus.UNHEALTHY
            
            status[service_type.value] = {
                'status': service_status.value,
                'available': self._is_service_available(service_type),
                'health_score': metrics.get_health_score(),
                'total_requests': metrics.total_requests,
                'success_rate': metrics.availability,
                'error_rate': metrics.error_rate,
                'average_response_time': metrics.average_response_time,
                'consecutive_failures': metrics.consecutive_failures,
                'consecutive_successes': metrics.consecutive_successes,
                'circuit_breaker_open': circuit_breaker.is_open,
                'last_success': metrics.last_success_time.isoformat() if metrics.last_success_time else None,
                'last_failure': metrics.last_failure_time.isoformat() if metrics.last_failure_time else None
            }
        
        return {
            'services': status,
            'preferred_service': self.preferred_service.value,
            'fallback_enabled': self.enable_fallback,
            'load_balancing_enabled': self.enable_load_balancing
        }
    
    def get_metrics_summary(self) -> Dict[str, Any]:
        """Get summary of service metrics."""
        total_requests = sum(m.total_requests for m in self.service_metrics.values())
        total_successes = sum(m.successful_requests for m in self.service_metrics.values())
        total_failures = sum(m.failed_requests for m in self.service_metrics.values())
        
        return {
            'total_requests': total_requests,
            'total_successes': total_successes,
            'total_failures': total_failures,
            'overall_success_rate': total_successes / total_requests if total_requests > 0 else 0.0,
            'service_distribution': {
                service_type.value: metrics.total_requests
                for service_type, metrics in self.service_metrics.items()
            }
        }
    
    async def shutdown(self):
        """Shutdown the router and cleanup resources."""
        if self._health_check_task:
            self._health_check_task.cancel()
            try:
                await self._health_check_task
            except asyncio.CancelledError:
                pass
        
        self.logger.info("Vision service router shutdown complete")


# Factory function
def create_vision_service_router(logger=None) -> VisionServiceRouter:
    """
    Create and configure a vision service router.
    
    Args:
        logger: Logger instance
        
    Returns:
        Configured VisionServiceRouter instance
    """
    return VisionServiceRouter(logger)